package io.sqrtqiezi.spark.explode

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

object ZipSample {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .master("local")
      .appName("zip sample")
      .getOrCreate()
    spark.sparkContext.setLogLevel("WARN")

    import spark.implicits._

    val data = Array(
      ("A", Array(1, 2), Array("a", "b")),
      ("B", Array(1, 2, 3), Array("a", "b", "c"))
    )

    val df = spark.createDataFrame(data)
      .toDF("key", "col1", "col2")
    df.show()

    df.createOrReplaceTempView("tbl")

//    df.select($"key", arrays_zip($"col1", $"col2").alias("zipped"))
//      .select($"key", explode($"zipped").alias("pair"))
//      .select($"key", $"pair.col1", $"pair.col2")
//      .show()

    spark.sql(
      """select key, pair.col1, pair.col2
        |from tbl lateral
        |view explode(arrays_zip(col1, col2)) as pair""".stripMargin)
      .show()

    spark.stop()
  }
}
